from yacf.Backends.C99.Visitors.LlcFilters import AccKernelsFilter
from yacf.Tools.Tree import NodeNotFound
from yacf.Backends.Common.Mutators.AbstractMutator import AbstractMutator

from yacf.Backends.Common.Visitors.GenericVisitors import GenericFilterVisitor
from yacf.Frontend.Omp import omp_ast

import re
from yacf.Tools.DescriptionFile.XMLParser import XMLManager

class LlcNameFilter(GenericFilterVisitor):
    """ Returns the first node with the given attribute
    """
    def __init__(self):
        def condition(node):
            if type(node) == omp_ast.llcName:
                return True
            return False
        super(LlcNameFilter, self).__init__(condition_func = condition)
    
class LlcForByNameFilter(GenericFilterVisitor):
    """ Returns the first llcFor node with the given attribute
    """
    def __init__(self, llcForName):
        def condition(node):
            if type(node) == omp_ast.llcFor:
                names = filter(lambda x: type(x) == omp_ast.llcName, node.clauses)
               # If more than one name clause appear, only the first is honored
                if names and (llcForName == names[0].label):
                       return True
            return False
        super(LlcForByNameFilter, self).__init__(condition_func = condition)

class XMLTranslationMutator(AbstractMutator):
    """ Translate yacf-frangollo llc for constructs with tranforms specified in `.xml` file.
        It applies to every llc region
    """

    def __init__(self, xml):
        assert isinstance(xml, XMLManager)
        self.xml = xml
        super(XMLTranslationMutator, self).__init__()


    def applyTransform(self, ast_llcFor, transformName, transformText):
        """ Select transform to apply to the llcFor loop
        """

        #===============================================================================
        # from yacf.Tools.Debug import DotDebugTool
        # DotDebugTool(highlight = []).apply(ast_llcFor)
        #===============================================================================   

        if (transformName == "tiling"):
            from yacf.MiddleEnd.Loop.Mutators.LoopTiling import TilingLoopMutator
            TilingLoopMutator(transformText).mutatorFunction(ast_llcFor)
        elif (transformName == "unswitch"):
            from yacf.MiddleEnd.Loop.Mutators.LoopUnSwitch import UnSwitchLoopMutator
            UnSwitchLoopMutator().mutatorFunction(ast_llcFor)
        else:
            raise NotImplemented

    def mutatorFunction(self, region):
        """ Llc mutator, for every llc for inside a single parallel region. Contrast the
            xml file with the ast, and call `self.applyTransform()` with each transform
            element found. 
            **Right now, it choose only the first cluster**
        """
        llcName = LlcNameFilter().apply(region)
        regionName = str(llcName.label).lstrip("\"").rstrip("\"")

        #===============================================================================
        # from yacf.Tools.Debug import DotDebugTool
        # DotDebugTool(highlight = [llcName,]).apply(region)
        #===============================================================================

        self.xml.prepare_iter_func()
        #self.xml.find_cluster("micluster")
        #print self.xml.cluster.get("name")

        self.xml.find_region(regionName)
        print "Region", self.xml.region.get("name"), "found."

        while (self.xml.next_element() is not None):
            elementName = self.xml.element.get("name")
            print "Searching " + elementName + " element ...",
            element = LlcForByNameFilter(elementName).apply(region)
            print elementName + " element found."
            print "Aplying transforms to " + elementName + "..."
            while (self.xml.next_transform() is not None):
                transformName = self.xml.transform.get("name")
                transformText = self.xml.transform.text
                print "\t" + transformName + " (" + transformText  + ")"
                self.applyTransform(element, transformName, transformText)

        #===============================================================================
        #print name.label
        #import pdb
        #pdb.set_trace()
        #===============================================================================


    def apply_all(self, ast):
        """ Apply mutation to all llc regions """
        #self.ast = ast
        f = AccKernelsFilter()
        try:
            for elem in f.iterate(ast):
                self.mutatorFunction(elem)
        except NodeNotFound as nf:
            print str(nf)
        except StopIteration:
            return None
        return ast


